package com.yfbao.horizon.common;

import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * 所有令牌生成周期时按照1s计算，
 * 即表示下面的限流表示每秒允许的请求数
 */
@Slf4j
public class RateLimiterUtil {

    private static final Integer DEFAULT_RATE = 20;
    private static final RateLimiter rateLimiter = RateLimiter.create(DEFAULT_RATE);
    private static final Map<String,RateLimiter> rateMap= new ConcurrentHashMap<>();

    public static void getRateLmt(){
        getRateLmt(1);
    }

    public static void getRateLmt(int permits){
        getRateLmt(permits,10,TimeUnit.MILLISECONDS);
    }

    public static void getRateLmt(int permits,long timewait,TimeUnit timeUnit){
        boolean b = rateLimiter.tryAcquire(permits,timewait, timeUnit);
        if(!b){
            log.warn("get rate limiter failed");
        }
    }


    /**
     * 自定义限流阈值的 限流获取
     * @param rateNum 自定义限流阈值
     */
    public static void getDefRateLmt(Integer rateNum){
        getDefRateLmt(rateNum,1);
    }

    /**
     * @param rateNum 自定义限流阈值
     * @param permits 需要获取的令牌数
     */
    public static void getDefRateLmt(Integer rateNum,Integer permits){
        getDefRateLmt(rateNum,permits,10);
    }

    /**
     *
     * @param rateNum 自定义限流阈值
     * @param permits 需要获取的令牌数
     * @param timeWait 获取令牌的等待时间
     */
    public static void getDefRateLmt(Integer rateNum,Integer permits,long timeWait){
        getDefRateLmt(rateNum,permits,timeWait,TimeUnit.MILLISECONDS);
    }


    /**
     *
     * @param rateNum 自定义限流阈值
     * @param permits 需要获取的令牌数
     * @param timewait 获取令牌的等待时间
     * @param timeUnit 等待时间单位
     */
    public static void getDefRateLmt(Integer rateNum,Integer permits,long timewait,TimeUnit timeUnit){
        RateLimiter rate = rateLimiter;
        if(rateNum != DEFAULT_RATE){
            String key = String.valueOf(rateNum);
            RateLimiter rateGet = rateMap.get(key);
            if(rateGet ==null){
               synchronized (rateMap){
                   RateLimiter rateReGet = rateMap.get(key);
                   if(rateReGet == null){
                       if(rateMap.size()>1000){
                           log.warn("rateLimter over 1000");
                       }
                       rateReGet = RateLimiter.create(rateNum);
                       rateMap.put(key,rateReGet);
                   }
                   rateGet = rateReGet;
               }
            }
            rate =  rateGet;
        }
        boolean b = rate.tryAcquire(permits,timewait, timeUnit);
        if(!b){
            log.warn("get rate limiter failed");
        }
    }


    public static void getDefLazyRateLmt(){
        getDefLazyRateLmt(1);
    }

    public static void getDefLazyRateLmt(Integer permits){
        getDefLazyRateLmt(permits,10,TimeUnit.MILLISECONDS);
    }

    public static void getDefLazyRateLmt(Integer permits, long timewait, TimeUnit timeUnit){
        getDefLazyRateLmt(DEFAULT_RATE,5,TimeUnit.SECONDS,permits,timewait,timeUnit);
    }

    /**
     * 获取启动延时限流
     * @param rateNum 限流阈值
     * @param warmupPeriod 启动延时长
     * @param unit 启动延时单位
     * @param permits 实时获取并发数
     * @param timewait 阻塞等待时间
     * @param timeUnit 等待时间单位
     */
    public static void getDefLazyRateLmt(Integer rateNum, long warmupPeriod, TimeUnit unit,
                                         Integer permits, long timewait, TimeUnit timeUnit) {
        String key = "lazy-"+String.valueOf(rateNum);
        RateLimiter rateGet = rateMap.get(key);
        if (rateGet == null) {
            synchronized (rateMap) {
                RateLimiter rateReGet = rateMap.get(key);
                if (rateReGet == null) {
                    if (rateMap.size() > 1000) {
                        log.warn("rateLimter over 1000");
                    }
                    rateReGet = RateLimiter.create(rateNum,warmupPeriod,unit);
                    rateMap.put(key, rateReGet);
                }
                rateGet = rateReGet;
            }
        }

        boolean b = rateGet.tryAcquire(permits, timewait, timeUnit);
        if (!b) {
            log.warn("get rate limiter failed");
        }
    }

}
